package cn.edu.recommender

import cn.edu.recommender.OfflineRecommender.MONGODB_RATING_COLLECTION
import org.apache.spark.mllib.recommendation.{ALS, MatrixFactorizationModel, Rating}
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.{Dataset, SparkSession}

import java.lang.Thread.sleep

object ALSTrainer {

  def main(args: Array[String]): Unit = {
    // 定义配置
    val config = Map(
      "spark.cores" -> "local[*]",
      "mongo.uri" -> "mongodb://localhost:27017/recommender",
      "mongo.db" -> "recommender"
    )

    val mongoConfig: MongoConfig = MongoConfig(config("mongo.uri"), config("mongo.db"))

    // 创建一个 SparkSession
    val spark: SparkSession = SparkSession.builder().appName("StatisticsRecommender").master(config("spark.cores")).getOrCreate()

    // 在对 DataFrame 和 Dataset 进行操作许多操作都需要这个包进行支持

    import spark.implicits._

    //数据加载进来
    val ratingRDD = spark
      .read
      .option("uri", mongoConfig.uri)
      .option("collection", MONGODB_RATING_COLLECTION)
      .format("com.mongodb.spark.sql")
      .load()
      .as[ProductRating]
      .map(rating => Rating(rating.userId, rating.productId, rating.score))
      .rdd
      .cache()


    // 将一个 RDD 随机切分成两个 RDD，用以划分训练集和测试集
    val splits = ratingRDD.randomSplit(Array(0.8, 0.2))
    val trainingRDD = splits(0)
    val testingRDD = splits(1)


    // 这里指定迭代次数为 5，rank 和 lambda 在几个值中选取调整
    val result = for (rank <- Array(50, 100, 200); lambda <- Array(1, 0.1, 0.01, 0.001))
      yield {
        val model = ALS.train(trainingRDD, rank, 5, lambda)
        val rmse = getRMSE(model, testingRDD)
        (rank, lambda, rmse)
      }

    // 按照 rmse 排序
    println()
    println(result.sortBy(_._3).mkString("\n"))


    sleep(300000)
    // 关闭 Spark
    spark.stop()
  }

  def getRMSE(model: MatrixFactorizationModel, testingRDD: RDD[Rating]): Double = {
    val input = testingRDD.map(item => (item.user, item.product))
    val output = model.predict(input)
    val real = testingRDD.map(item => ((item.user, item.product), item.rating))
    val predict = output.map(item => ((item.user, item.product), item.rating))
    // 计算 RMSE
    math.sqrt(
      real.join(predict)
        .map {
          case ((userId, productId), (real, pre)) =>
            // 真实值和预测值之间的差
            val err = real - pre
            err * err
        }
        .mean()
    )
  }
}
